#!/usr/bin/env python
#  -*- coding: utf-8 -*-
"""
交易日历处理
"""
__author__ = 'mayanqiong'

import os
from typing import Union, List

import requests
from datetime import date, datetime

import pandas as pd


# 2003-01-01 ～ 2021-12-31 节假日，不包含周末
CHINESE_REST_DAYS = ['2003-01-01', '2003-01-30', '2003-01-31', '2003-02-03', '2003-02-04', '2003-02-05', '2003-02-06', '2003-02-07', '2003-05-01', '2003-05-02', '2003-05-05', '2003-05-06', '2003-05-07', '2003-10-01', '2003-10-02', '2003-10-03', '2003-10-06', '2003-10-07', '2004-01-01', '2004-01-19', '2004-01-20', '2004-01-21', '2004-01-22', '2004-01-23', '2004-01-26', '2004-01-27', '2004-01-28', '2004-05-03', '2004-05-04', '2004-05-05', '2004-05-06', '2004-05-07', '2004-10-01', '2004-10-04', '2004-10-05', '2004-10-06', '2004-10-07', '2005-01-03', '2005-02-07', '2005-02-08', '2005-02-09', '2005-02-10', '2005-02-11', '2005-02-14', '2005-02-15', '2005-05-02', '2005-05-03', '2005-05-04', '2005-05-05', '2005-05-06', '2005-10-03', '2005-10-04', '2005-10-05', '2005-10-06', '2005-10-07', '2006-01-02', '2006-01-03', '2006-01-30', '2006-01-31', '2006-02-01', '2006-02-02', '2006-02-03', '2006-05-01', '2006-05-02', '2006-05-03', '2006-05-04', '2006-05-05', '2006-10-02', '2006-10-03', '2006-10-04', '2006-10-05', '2006-10-06', '2007-01-01', '2007-01-02', '2007-01-03', '2007-02-19', '2007-02-20', '2007-02-21', '2007-02-22', '2007-02-23', '2007-05-01', '2007-05-02', '2007-05-03', '2007-05-04', '2007-05-07', '2007-10-01', '2007-10-02', '2007-10-03', '2007-10-04', '2007-10-05', '2007-12-31', '2008-01-01', '2008-02-06', '2008-02-07', '2008-02-08', '2008-02-11', '2008-02-12', '2008-04-04', '2008-05-01', '2008-05-02', '2008-06-09', '2008-09-15', '2008-09-29', '2008-09-30', '2008-10-01', '2008-10-02', '2008-10-03', '2009-01-01', '2009-01-02', '2009-01-26', '2009-01-27', '2009-01-28', '2009-01-29', '2009-01-30', '2009-04-06', '2009-05-01', '2009-05-28', '2009-05-29', '2009-10-01', '2009-10-02', '2009-10-05', '2009-10-06', '2009-10-07', '2009-10-08', '2010-01-01', '2010-02-15', '2010-02-16', '2010-02-17', '2010-02-18', '2010-02-19', '2010-04-05', '2010-05-03', '2010-06-14', '2010-06-15', '2010-06-16', '2010-09-22', '2010-09-23', '2010-09-24', '2010-10-01', '2010-10-04', '2010-10-05', '2010-10-06', '2010-10-07', '2011-01-03', '2011-02-02', '2011-02-03', '2011-02-04', '2011-02-07', '2011-02-08', '2011-04-04', '2011-04-05', '2011-05-02', '2011-06-06', '2011-09-12', '2011-10-03', '2011-10-04', '2011-10-05', '2011-10-06', '2011-10-07', '2012-01-02', '2012-01-03', '2012-01-23', '2012-01-24', '2012-01-25', '2012-01-26', '2012-01-27', '2012-04-02', '2012-04-03', '2012-04-04', '2012-04-30', '2012-05-01', '2012-06-22', '2012-10-01', '2012-10-02', '2012-10-03', '2012-10-04', '2012-10-05', '2013-01-01', '2013-01-02', '2013-01-03', '2013-02-11', '2013-02-12', '2013-02-13', '2013-02-14', '2013-02-15', '2013-04-04', '2013-04-05', '2013-04-29', '2013-04-30', '2013-05-01', '2013-06-10', '2013-06-11', '2013-06-12', '2013-09-19', '2013-09-20', '2013-10-01', '2013-10-02', '2013-10-03', '2013-10-04', '2013-10-07', '2014-01-01', '2014-01-31', '2014-02-03', '2014-02-04', '2014-02-05', '2014-02-06', '2014-04-07', '2014-05-01', '2014-05-02', '2014-06-02', '2014-09-08', '2014-10-01', '2014-10-02', '2014-10-03', '2014-10-06', '2014-10-07', '2015-01-01', '2015-01-02', '2015-02-18', '2015-02-19', '2015-02-20', '2015-02-23', '2015-02-24', '2015-04-06', '2015-05-01', '2015-06-22', '2015-09-03', '2015-09-04', '2015-10-01', '2015-10-02', '2015-10-05', '2015-10-06', '2015-10-07', '2016-01-01', '2016-02-08', '2016-02-09', '2016-02-10', '2016-02-11', '2016-02-12', '2016-04-04', '2016-05-02', '2016-06-09', '2016-06-10', '2016-09-15', '2016-09-16', '2016-10-03', '2016-10-04', '2016-10-05', '2016-10-06', '2016-10-07', '2017-01-02', '2017-01-27', '2017-01-30', '2017-01-31', '2017-02-01', '2017-02-02', '2017-04-03', '2017-04-04', '2017-05-01', '2017-05-29', '2017-05-30', '2017-10-02', '2017-10-03', '2017-10-04', '2017-10-05', '2017-10-06', '2018-01-01', '2018-02-15', '2018-02-16', '2018-02-19', '2018-02-20', '2018-02-21', '2018-04-05', '2018-04-06', '2018-04-30', '2018-05-01', '2018-06-18', '2018-09-24', '2018-10-01', '2018-10-02', '2018-10-03', '2018-10-04', '2018-10-05', '2018-12-31', '2019-01-01', '2019-02-04', '2019-02-05', '2019-02-06', '2019-02-07', '2019-02-08', '2019-04-05', '2019-05-01', '2019-05-02', '2019-05-03', '2019-06-07', '2019-09-13', '2019-10-01', '2019-10-02', '2019-10-03', '2019-10-04', '2019-10-07', '2020-01-01', '2020-01-24', '2020-01-27', '2020-01-28', '2020-01-29', '2020-01-30', '2020-01-31', '2020-04-06', '2020-05-01', '2020-05-04', '2020-05-05', '2020-06-25', '2020-06-26', '2020-10-01', '2020-10-02', '2020-10-05', '2020-10-06', '2020-10-07', '2020-10-08', '2021-01-01', '2021-02-11', '2021-02-12', '2021-02-15', '2021-02-16', '2021-02-17', '2021-04-05', '2021-05-03', '2021-05-04', '2021-05-05', '2021-06-14', '2021-09-20', '2021-09-21', '2021-10-01', '2021-10-04', '2021-10-05', '2021-10-06', '2021-10-07']
rest_days_df = pd.DataFrame(data={'date': pd.Series(pd.to_datetime(CHINESE_REST_DAYS, format='%Y-%m-%d'))})
rest_days_df['trading_restdays'] = False  # 节假日为 False


def _get_trading_calendar(start_dt: date, end_dt: date):
    """
    获取一段时间内，每天是否是交易日

    :return: DataFrame
        date  trading
    2019-12-05  True
    2019-12-06  True
    2019-12-07  False
    2019-12-08  False
    2019-12-09  True
    """
    df = pd.DataFrame()
    df['date'] = pd.Series(pd.date_range(start=start_dt, end=end_dt, freq="D"))
    df['trading'] = df['date'].dt.dayofweek.lt(5)
    result = pd.merge(rest_days_df, df, sort=True, how="right", on="date")
    result.fillna(True, inplace=True)
    df['trading'] = result['trading'] & result['trading_restdays']
    return df


class TqContCalendar(object):

    """
    主连日历
    df:
        date   trading  KQ.m@DCE.a  KQ.m@DCE.eg
    2019-12-06    True  DCE.a2005  DCE.eg2001
    2019-12-09    True  DCE.a2005  DCE.eg2001
    2019-12-10    True  DCE.a2005  DCE.eg2005
    2019-12-11    True  DCE.a2005  DCE.eg2005
    2019-12-12    True  DCE.a2005  DCE.eg2005
    """

    continuous = None

    def __init__(self, start_dt: date, end_dt: date, symbols: Union[List[str], None] = None) -> None:
        """
        初始化主连日历表
        :param date start_dt: 开始交易日日期
        :param date end_dt: 结束交易日日期
        :param list[str] symbols: 主连合约列表
        :return:
        """
        self.df = _get_trading_calendar(start_dt=start_dt, end_dt=end_dt)
        self.df = self.df.loc[self.df.trading, ['date']]  # 只保留交易日
        self.df.reset_index(inplace=True, drop=True)
        if TqContCalendar.continuous is None:
            rsp = requests.get(os.getenv("TQ_CONT_TABLE_URL", "https://files.shinnytech.com/continuous_table.json"))  # 下载历史主连合约信息
            rsp.raise_for_status()
            TqContCalendar.continuous = {f"KQ.m@{k}": v for k, v in rsp.json().items()}
        if symbols is not None:
            if not all([s in TqContCalendar.continuous.keys() for s in symbols]):
                raise Exception(f"参数错误，symbols={symbols} 中应该全部都是主连合约代码")
        symbols = TqContCalendar.continuous.keys() if symbols is None else symbols
        self.start_dt, self.end_dt = self.df.iloc[0].date, self.df.iloc[-1].date
        for s in symbols:
            self._ensure_cont_on_df(s)

    def _ensure_cont_on_df(self, cont):
        """将一个主连对应的标的填在 self.df 对应位置"""
        temp_df = pd.DataFrame(data=TqContCalendar.continuous[cont], columns=['date', 'underlying'])
        temp_df['date'] = pd.Series(pd.to_datetime(temp_df['date'], format='%Y%m%d'))
        merge_result = pd.merge(temp_df, self.df, sort=True, how="outer", on="date")
        merge_result.fillna(method="ffill", inplace=True)
        merge_result.fillna(value="", inplace=True)
        s = merge_result.loc[merge_result.date.ge(self.start_dt) & merge_result.date.le(self.end_dt), 'underlying']
        self.df[cont] = pd.Series(s.values)

    def _get_cont_underlying_on_date(self, dt: datetime):
        """返回某一交易日的全部主连"""
        df = self.df.loc[self.df.date.ge(dt), :]
        return df.iloc[0:1]
